"""
    @brief:基于互补滤波(Complimentary Filter)算法的6轴IMU姿态解算
    @author:HuYunhao
    @email:hnu_huyunhao1995@163.com
    @reference:http://philsal.co.uk/projects/imu-attitude-estimation
"""
import numpy as np
from math import sin, cos, tan, pi


class ComplimentaryFilter:
    def __init__(self):
        self.dt = 0.01  # IMU的采样周期
        self.alpha = 0.02  # 加速度计的置信度
        self.roll_hat = 0  # 使用互补滤波算法融合后的roll估计值
        self.pitch_hat = 0
        self.yaw_hat = 0
        self.gyro_angle_velocity = []   # 大地坐标系下的姿态角速度

    def get_gyro_attitude_estimate(self, gyro_data: list):
        """
        @brief 通过陀螺仪的输入[gx, gy, gz]，计算大地坐标系下的姿态角速度[roll', pitch', yaw']
        """
        r = self.roll_hat
        p = self.pitch_hat
        y = self.yaw_hat
        trans = np.array([[1, sin(r) * tan(p), cos(r) * tan(p)],  # 由陀螺仪输出到姿态角速度的转换矩阵
                          [0, cos(r), -sin(r)],
                          [0, sin(r) / cos(p), cos(r) / cos(p)]])
        gyro = np.array(gyro_data).reshape(3, 1)
        self.gyro_angle_velocity = np.matmul(trans, gyro).reshape(3).tolist()

    def complimentary_filter(self, acc_hat: list, gyro_data: list):
        """
        @brief 融合加速度计和陀螺仪的估计值
        """
        self.get_gyro_attitude_estimate(gyro_data)
        a = self.alpha
        r = self.roll_hat
        p = self.pitch_hat
        self.roll_hat = (1 - a) * (r + self.gyro_angle_velocity[0] * self.dt) + a * acc_hat[0]
        self.pitch_hat = (1 - a) * (p + self.gyro_angle_velocity[1] * self.dt) + a * acc_hat[1]
        self.yaw_hat = self.yaw_hat + self.gyro_angle_velocity[2] * self.dt
        return [self.roll_hat, self.pitch_hat, self.yaw_hat]
